import os
import subprocess
from macholib.MachO import MachO
from graphviz import Digraph

def is_mach_o_file(file_path):
    try:
        macho = MachO(file_path)
        return True
    except:
        return False

def get_dependencies(file_path):
    # print(f"正在获取文件依赖: {file_path}")
    macho = MachO(file_path)
    dependencies = []
    # 定义加载命令常量
    LC_LOAD_DYLIB = 0x0C  # 十进制12
    LC_LOAD_WEAK_DYLIB = 0x80000018  # 十进制2147483672
    
    for header in macho.headers:
        for cmd in header.commands:
            cmd_type = cmd[0].cmd
            if cmd_type in (LC_LOAD_DYLIB, LC_LOAD_WEAK_DYLIB):
                # 提取依赖库的名称
                dylib_path = cmd[2].decode('utf-8')
                dependencies.append(dylib_path)

    # print(f"文件依赖: {dependencies}")                
    return dependencies

# 添加树形结构相关的类
class DependencyNode:
    def __init__(self, name, path=None):
        self.name = name
        self.path = path
        self.children = []
        self.visited = False  # 用于检测循环依赖
    
    def add_child(self, child_node):
        self.children.append(child_node)
    
    def __str__(self):
        return self.name

def find_all_mach_o_files(directory):
    """递归查找目录下所有Mach-O文件"""
    mach_o_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            if is_mach_o_file(file_path):
                mach_o_files.append(file_path)
    return mach_o_files

def find_unused_libraries(all_mach_o_files, tree_root=None):
    """查找未被依赖的动态库
    
    Args:
        all_mach_o_files (list): 所有Mach-O文件的路径列表
        tree_root (DependencyNode, optional): 依赖树的根节点
        
    Returns:
        list: 未被依赖的库列表
    """
    used_libraries = set()
    all_libraries = set()
    
    # 首先收集所有库
    for file_path in all_mach_o_files:
        all_libraries.add(os.path.basename(file_path))
    
    # 如果有依赖树根节点，从树中收集所有被依赖的库
    if tree_root:
        # 递归遍历依赖树，收集所有节点
        def collect_used_libraries(node):
            if node.path:
                used_libraries.add(os.path.basename(node.path))
            for child in node.children:
                collect_used_libraries(child)
        
        # 从根节点开始收集
        collect_used_libraries(tree_root)
    else:
        # 如果没有依赖树，使用原来的方法收集被依赖的库
        for file_path in all_mach_o_files:
            try:
                deps = get_dependencies(file_path)
                for dep in deps:
                    dep_name = os.path.basename(dep)
                    # 检查这个依赖是否存在于所有库中
                    for lib in all_libraries:
                        if lib in dep_name:
                            used_libraries.add(lib)
            except Exception as e:
                print(f"分析文件依赖时出错 {file_path}: {e}")
    
    print(f"所有库: {all_libraries},\n已用库: {used_libraries}")
    # 找出未被依赖的库
    unused_libraries = all_libraries - used_libraries
    return sorted(list(unused_libraries))  # 排序以保持稳定的显示顺序

def analyze_directory(directory, main_executable=None):
    # 检查目录是否存在
    if not os.path.exists(directory):
        print(f"错误: 目录 {directory} 不存在")
        return Digraph(comment='Mach-O Dependencies')
    
    if not os.path.isdir(directory):
        print(f"错误: {directory} 不是有效目录")
        return Digraph(comment='Mach-O Dependencies')

    # 创建有向图
    dependency_graph = Digraph(comment='Mach-O Dependencies')
    dependency_graph.attr(rankdir='TB')
    dependency_graph.attr('node', shape='box')
    
    # 查找所有Mach-O文件
    all_mach_o_files = find_all_mach_o_files(directory)
    
    # 预先获取所有Mach-O文件的依赖关系
    dependency_cache = {}
    print("预先获取所有Mach-O文件的依赖关系...")
    for file_path in all_mach_o_files:
        try:
            file_name = os.path.basename(file_path)
            dependency_cache[file_path] = get_dependencies(file_path)

        except Exception as e:
            print(f"获取 {file_path} 依赖时出错: {e}")
    
    if not main_executable:
        # 如果没有指定主程序，显示所有Mach-O文件及其直接依赖
        dependency_map = {}
        for file_path in all_mach_o_files:
            file_name = os.path.basename(file_path)
            dependency_map[file_name] = []
            try:
                # 使用缓存的依赖关系
                dependencies = dependency_cache[file_path]
                for dep in dependencies:
                    dep_name = os.path.basename(dep)
                    dependency_map[file_name].append(dep_name)
            except Exception as e:
                print(f"处理文件 {file_path} 时出错: {e}")
        
        # 为每个Mach-O文件创建节点
        for file_name in dependency_map:
            with dependency_graph.subgraph(name=f"cluster_{file_name}") as subg:
                subg.attr(label=file_name, style='filled', fillcolor='lightyellow')
                subg.node(file_name, shape='box', style='filled', fillcolor='lightgreen')
                
                # 添加依赖边
                for dep in dependency_map[file_name]:
                    dep_node = f"{file_name}_{dep}"
                    subg.node(dep_node, label=dep, shape='box', style='filled', fillcolor='lightblue')
                    subg.edge(file_name, dep_node)
    else:
        # 以主程序为根构建依赖树
        if not os.path.isabs(main_executable):
            main_executable = os.path.join(directory, main_executable)
        
        if main_executable not in all_mach_o_files:
            print(f"错误: 主程序 {main_executable} 不是Mach-O文件或不在指定目录中")
            return dependency_graph
        
        # 使用缓存的依赖关系构建树，并传入lib_depths字典
        tree_root = build_dependency_tree_with_cache(main_executable, all_mach_o_files, dependency_cache, lib_depths={})
        render_tree_to_graphviz(tree_root, dependency_graph)
    
    return dependency_graph

def build_dependency_tree_with_cache(main_executable, all_mach_o_files, dependency_cache, visited=None, depth=0, max_depth=20, lib_depths=None):
    """使用缓存的依赖关系构建依赖树"""
    if not os.path.exists(main_executable):
        print(f"错误: 主程序文件 {main_executable} 不存在")
        return DependencyNode(os.path.basename(main_executable))
    
    if visited is None:
        visited = set()
    
    if lib_depths is None:
        lib_depths = {}  # 记录每个库的最大深度
    
    if depth > max_depth:
        print(f"达到最大递归深度 {max_depth}，停止构建")
        return DependencyNode(os.path.basename(main_executable))
    
    node = DependencyNode(os.path.basename(main_executable), main_executable)
    visited.add(main_executable)
    
    try:
        # 使用缓存的依赖关系
        if main_executable in dependency_cache:
            dependencies = dependency_cache[main_executable]
            if "Aweme" in os.path.basename(main_executable):
                print(f"{main_executable} 的依赖关系: {dependencies}")

            for dep in dependencies:
                dep_path = None
                dep_name = os.path.basename(dep)


                for mach_o_file in all_mach_o_files:
                    # 使用rstrip()移除末尾的空字符
                    clean_dep_name = dep_name.rstrip('\x00').strip()
                    clean_file_name = os.path.basename(mach_o_file).strip()
                    
                    if "AwemeCore" in clean_dep_name:
                        print(f"清理后比较: dep_name='{clean_dep_name}'({len(clean_dep_name)}), file_name='{clean_file_name}'({len(clean_file_name)})\n")
                    
                    if clean_dep_name == clean_file_name:
                        dep_path = mach_o_file
                        break
                print(f"dep_name==={dep_name},dep_path={dep_path}\n")
                if dep_path:
                    current_depth = depth + 1
                    lib_name = os.path.basename(dep_path)
                    
                    # 检查是否应该在当前深度处理这个依赖
                    if lib_name in lib_depths:
                        if current_depth <= lib_depths[lib_name]:
                            # 跳过当前节点，因为在更深的层级已经或将要处理
                            continue
                    
                    # 更新依赖的最大深度
                    lib_depths[lib_name] = current_depth

                    if "libswiftCoreGraphics" in dep_name :
                        print(f"dep_name=libswiftCoreGraphics,\ndep_path={dep_path}\n,dep_name={dep_name},\nvisited={visited}\n")
                    
                    if dep_path not in visited:
                        child_node = build_dependency_tree_with_cache(
                            dep_path, 
                            all_mach_o_files, 
                            dependency_cache,
                            visited.copy(), 
                            current_depth, 
                            max_depth,
                            lib_depths
                        )
                        print(f"child_node={child_node},\n dep_name={dep_name},\nvisited={visited},\ndep_path={dep_path}\n")
                        node.add_child(child_node)
        else:
            print(f"警告: 未找到 {main_executable} 的缓存依赖")
    except Exception as e:
        print(f"处理文件 {main_executable} 时出错: {e}")
    
    return node

def render_tree_to_graphviz(tree_node, graph, parent=None, depth=0, path=""):
    """将树渲染为Graphviz图"""
    # 生成节点ID，确保唯一性
    node_id = f"{path}_{tree_node.name}" if path else tree_node.name
    node_id = node_id.replace(".", "_")  # 替换特殊字符
    
    # 设置节点样式
    node_attrs = {
        "label": tree_node.name,
        "shape": "box",
        "style": "filled",
        "fillcolor": "lightgreen" if depth == 0 else "lightblue"
    }
    
    # 添加节点
    graph.node(node_id, **node_attrs)
    
    # 添加边
    if parent:
        graph.edge(parent, node_id)
    
    # 递归处理子节点
    for i, child in enumerate(tree_node.children):
        child_path = f"{path}_{i}" if path else str(i)
        render_tree_to_graphviz(child, graph, node_id, depth + 1, child_path)

def save_graph_to_html(graph, output_path, directory_path=None, tree_root=None):
    # 确保输出目录可写
    output_dir = os.path.dirname(os.path.abspath(output_path)) or os.getcwd()
    try:
        os.makedirs(output_dir, exist_ok=True)
        if not os.access(output_dir, os.W_OK):
            raise PermissionError(f"无写入权限: {output_dir}")
    except (PermissionError, OSError) as e:
        print(f"警告: {e}, 使用临时目录替代")
        import tempfile
        output_dir = tempfile.gettempdir()
        output_path = os.path.join(output_dir, os.path.basename(output_path))
    
    # 从输出路径中提取应用名称
    app_name = os.path.basename(output_path).replace('_dependency.html', '')
    graph_name = app_name  # 使用应用名称作为文件名
    
    # 使用绝对路径
    output_dir = os.path.dirname(os.path.abspath(output_path))
    graph_path = os.path.join(output_dir, graph_name)
    
    # 保存并渲染图
    graph.save(f"{graph_path}.gv")
    try:
        # 渲染为SVG格式
        graph.render(graph_path, format='svg', view=False, cleanup=True)
        print(f"SVG图像已保存到 {graph_path}.svg")
        
        # 设置img_path为相对路径
        img_path = f"./{graph_name}.svg"
    except Exception as e:
        print(f"渲染图像时出错: {e}")
        return
    
    # 获取未使用的库
    unused_libs = []
    if directory_path:
        all_mach_o_files = find_all_mach_o_files(directory_path)
        unused_libs = find_unused_libraries(all_mach_o_files, tree_root)
        
        # 统计所有动态库的出现次数
        lib_counts = {}
        system_lib_paths = ['/usr/lib/', '/System/Library/', '/Library/']
        for file_path in all_mach_o_files:
            try:
                deps = get_dependencies(file_path)
                for dep in deps:
                    dep_name = os.path.basename(dep)
                    # 过滤系统库和libswift开头的库
                    if (not any(dep.startswith(sys_path) for sys_path in system_lib_paths) 
                        and not dep_name.startswith('libswift')):
                        lib_counts[dep_name] = lib_counts.get(dep_name, 0) + 1
            except Exception as e:
                print(f"统计依赖时出错 {file_path}: {e}")
            
            # 按出现次数排序
            sorted_lib_counts = sorted(lib_counts.items(), key=lambda x: x[1], reverse=True)
    
    # 生成表格HTML
    unused_table_html = ""
    if unused_libs:
        unused_table_html = """
            <div style="width: 100%; margin: 20px 0;">
                <h2>未被依赖的动态库</h2>
                <table class="unused-table">
                    <thead>
                        <tr>
                            <th>序号</th>
                            <th>库名称</th>
                        </tr>
                    </thead>
                    <tbody>
                        {rows}
                    </tbody>
                </table>
            </div>
        """.format(rows=''.join(f"<tr><td>{i+1}</td><td>{lib}</td></tr>" for i, lib in enumerate(unused_libs)))
    
    # 生成依赖统计表格HTML
    lib_count_table_html = ""
    if directory_path and lib_counts:
        lib_count_table_html = """
            <div style="width: 100%; margin: 20px 0;">
                <h2>非系统动态库依赖统计</h2>
                <table class="lib-count-table">
                    <thead>
                        <tr>
                            <th>序号</th>
                            <th>库名称</th>
                            <th>被引用次数</th>
                        </tr>
                    </thead>
                    <tbody>
                        {rows}
                    </tbody>
                </table>
            </div>
        """.format(rows=''.join(
            f"<tr><td>{i+1}</td><td>{lib}</td><td>{count}</td></tr>" 
            for i, (lib, count) in enumerate(sorted_lib_counts)
        ))
    
    # 在HTML模板中添加lib_count_table_html
    html_content = f"""
    <!DOCTYPE html>
    <html lang="zh-CN">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Mach-O 依赖树</title>
        <style>
            body {{
                font-family: Arial, sans-serif;
                margin: 20px;
            }}
            h1 {{
                color: #333;
                text-align: center;
            }}
            .container {{
                max-width: 1200px;
                margin: 0 auto;
                padding: 20px;
            }}
            .svg-container {{
                width: 90%;
                height: 70vh;
                overflow: auto;
                border: 1px solid #ddd;
                margin: 20px auto;
                padding: 20px;
                box-shadow: 0 0 10px rgba(0,0,0,0.1);
            }}
            .svg-container object {{
                width: 100%;
                height: 100%;
            }}
            .legend {{
                display: flex;
                flex-wrap: wrap;
                margin: 10px 0;
                padding: 10px;
                background-color: #f9f9f9;
                border-radius: 5px;
            }}
            .legend-item {{
                margin-right: 20px;
            }}
            .green {{
                color: green;
                font-weight: bold;
            }}
            .blue {{
                color: blue;
                font-weight: bold;
            }}
            .unused-table, .lib-count-table {{
                margin-top: 20px;
                width: 100%;
                border-collapse: collapse;
                background-color: #fff;
            }}
            .unused-table th, .unused-table td,
            .lib-count-table th, .lib-count-table td {{
                border: 1px solid #ddd;
                padding: 8px;
                text-align: left;
            }}
            .unused-table th, .lib-count-table th {{
                background-color: #f5f5f5;
            }}
            .unused-table tr:nth-child(even),
            .lib-count-table tr:nth-child(even) {{
                background-color: #f9f9f9;
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>Mach-O 依赖多叉树</h1>
            
            <div class="legend">
                <h3>图例说明:</h3>
                <div class="legend-item"><span class="green">绿色节点</span>: 主要Mach-O文件（根节点）</div>
                <div class="legend-item"><span class="blue">蓝色节点</span>: 依赖库（子节点）</div>
            </div>
            
            <div class="svg-container">
                <object data="{img_path}" type="image/svg+xml" oncontextmenu="window.open(this.data, '_blank'); return false;"></object>
            </div>
            
            {lib_count_table_html}
            {unused_table_html}
        </div>
    </body>
    </html>
    """
    
    # 删除下面重复的HTML模板定义
    # 保留上面完整的HTML模板定义
    
    # 写入HTML文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(html_content)
    print(f"HTML文件已保存到 {output_path}")
    
    # 在浏览器中打开文件
    try:
        import webbrowser
        webbrowser.open(f'file://{os.path.abspath(output_path)}')
        print(f"已在浏览器中打开 {output_path}")
    except Exception as e:
        print(f"无法在浏览器中打开文件: {e}")

def analyze_mach_o_dependencies(directory_path, output_path, main_executable=None):
    """
    分析Mach-O文件依赖关系并生成可视化结果
    
    Args:
        directory_path (str): 要分析的目录路径
        output_path (str): 输出HTML文件的路径
        main_executable (str, optional): 主程序文件名（相对于directory_path）
    
    Returns:
        bool: 分析是否成功
    """
    try:
        directory_path = os.path.abspath(directory_path)
        output_path = os.path.abspath(output_path)
        
        print(f"开始分析目录: {directory_path}")
        
        dependency_graph = analyze_directory(directory_path, main_executable)
        
        # 获取树根节点（如果有主程序）
        tree_root = None
        if main_executable:
            all_mach_o_files = find_all_mach_o_files(directory_path)
            dependency_cache = {}
            for file_path in all_mach_o_files:
                try:
                    dependency_cache[file_path] = get_dependencies(file_path)
                except Exception:
                    pass
            
            main_path = os.path.join(directory_path, main_executable) if not os.path.isabs(main_executable) else main_executable
            if main_path in all_mach_o_files:
                tree_root = build_dependency_tree_with_cache(main_path, all_mach_o_files, dependency_cache, lib_depths={})
        
        # 传递directory_path和tree_root给save_graph_to_html
        save_graph_to_html(dependency_graph, output_path, directory_path, tree_root)
        
        return True
        
    except Exception as e:
        print(f"分析过程中出错: {e}")
        return False

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Analyze Mach-O dependencies')
    parser.add_argument('directory', type=str, help='Directory to analyze')
    parser.add_argument('output', type=str, help='Output HTML file')
    parser.add_argument('--main', type=str, help='Main executable file (relative to directory)', default=None)
    args = parser.parse_args()

    # 使用封装函数
    analyze_mach_o_dependencies(args.directory, args.output, args.main)